#include "tensor_mat_multi_cav.h"
#include <cmath>
#include <iostream>
#include <vector>
#include <algorithm>
#include <numeric>
#include <cassert>
#include <mpi.h>

namespace TensorMatMultiCavity {

// Legendre-Gauss-Lobatto points and weights
void legendre_gauss_lobatto(int N, std::vector<double>& points, std::vector<double>& weights) {
    // Compute interior points and weights using Jacobi-Gauss quadrature
    std::vector<double> interior_points(N - 2);
    std::vector<double> interior_weights(N - 2);
    
    // Use Jacobi-Gauss quadrature for interior points
    TensorMatMultiCavity::jacobi_gauss_quadrature(N - 2, 1, 1, interior_points, interior_weights);
    
    // Add boundary points
    points.resize(N);
    weights.resize(N);
    points[0] = -1.0;
    points[N - 1] = 1.0;
    
    // Copy interior points
    for (int i = 0; i < N - 2; i++) {
        points[i + 1] = interior_points[i];
    }
    
    // Compute weights using Legendre polynomial
    for (int i = 0; i < N; i++) {
        double L_n = TensorMatMultiCavity::legendre_polynomial(N, points[i]);
        weights[i] = 2.0 / (N * (N - 1) * L_n * L_n);
    }
}

// Jacobi-Gauss quadrature points and weights
void jacobi_gauss_quadrature(int N, double alpha, double beta, std::vector<double>& points, std::vector<double>& weights) {
    // Compute recurrence coefficients
    std::vector<double> alpha_vec(N);
    std::vector<double> beta_vec(N);
    
    for (int i = 0; i < N; i++) {
        double a = alpha + i;
        double b = beta + i;
        double ab = a + b;
        
        alpha_vec[i] = (b * b - a * a) / ((ab + 2) * ab);
        if (i == 0) {
            beta_vec[i] = std::pow(2, ab + 1) * std::tgamma(a + 1) * std::tgamma(b + 1) / std::tgamma(ab + 2);
        } else {
            beta_vec[i] = 4 * (i + a) * (i + b) * i * (i + ab) / ((ab + 2 * i) * (ab + 2 * i + 1) * (ab + 2 * i - 1));
        }
    }
    
    // Solve eigenvalue problem
    TensorMatMultiCavity::eigen_tridiagonal(alpha_vec, beta_vec, points, weights);
}

// Legendre polynomial evaluation
double legendre_polynomial(int n, double x) {
    if (n == 0) return 1.0;
    if (n == 1) return x;
    
    double p0 = 1.0;
    double p1 = x;
    double pn = 0.0;
    
    for (int i = 2; i <= n; i++) {
        pn = ((2 * i - 1) * x * p1 - (i - 1) * p0) / i;
        p0 = p1;
        p1 = pn;
    }
    
    return pn;
}

// Placeholder for eigenvalue computation of tridiagonal matrix
// In a real implementation, use LAPACK or optimized libraries
void eigen_tridiagonal(const std::vector<double>& alpha, const std::vector<double>& beta,
                      std::vector<double>& eigenvalues, std::vector<double>& eigenvectors) {
    // Implementation needed - this is a placeholder
    // Should use a robust eigenvalue solver for tridiagonal matrices
    // For now, just set some default values
    int N = alpha.size();
    eigenvalues.resize(N);
    eigenvectors.resize(N * N);
    
    for (int i = 0; i < N; i++) {
        eigenvalues[i] = alpha[i];
        for (int j = 0; j < N; j++) {
            eigenvectors[i * N + j] = (i == j) ? 1.0 : 0.0;
        }
    }
}

// Compute derivative matrix for Legendre-Gauss-Lobatto points
void differentiation_matrix(const std::vector<double>& points, std::vector<std::vector<double>>& D) {
    int N = points.size() - 1;
    D.resize(N + 1, std::vector<double>(N + 1, 0.0));
    
    // Compute barycentric weights
    std::vector<double> w(N + 1, 1.0);
    for (int i = 0; i <= N; i++) {
        for (int j = 0; j <= N; j++) {
            if (i != j) {
                w[i] *= (points[i] - points[j]);
            }
        }
        w[i] = 1.0 / w[i];
    }
    
    // Compute differentiation matrix
    for (int i = 0; i <= N; i++) {
        for (int j = 0; j <= N; j++) {
            if (i != j) {
                D[i][j] = w[j] / (w[i] * (points[i] - points[j]));
            }
        }
    }
    
    // Compute diagonal entries
    for (int i = 0; i <= N; i++) {
        double sum = 0.0;
        for (int j = 0; j <= N; j++) {
            if (i != j) {
                sum += D[i][j];
            }
        }
        D[i][i] = -sum;
    }
}

// Generate SEM operators for 1D
void generate_sem_operators_1d(int N, int Ncell, double L, 
                              std::vector<double>& points, 
                              std::vector<std::vector<double>>& D1, 
                              std::vector<std::vector<double>>& D2) {
    // Ensure minimum values for N and Ncell
    N = std::max(1, N);
    Ncell = std::max(1, Ncell);
    
    std::cout << "SEM Operator Generation: N=" << N << ", Ncell=" << Ncell << ", L=" << L << std::endl;
    
    // Generate Legendre-Gauss-Lobatto points and weights for a single element
    std::vector<double> element_points;
    std::vector<double> weights;
    legendre_gauss_lobatto(N, element_points, weights);
    
    // Map the points to the physical domain
    double dx = 2.0 * L / Ncell;
    points.clear();
    
    for (int cell = 0; cell < Ncell; cell++) {
        double cell_left = -L + cell * dx;
        for (int i = 0; i <= N; i++) {
            // Only add the first point of each element except for the first element
            if (i == 0 && cell > 0) continue;
            
            double x = cell_left + dx/2.0 * (element_points[i] + 1.0);
            points.push_back(x);
        }
    }
    
    // Compute differentiation matrix for a single element
    std::vector<std::vector<double>> element_D;
    differentiation_matrix(element_points, element_D);
    
    // Scale by element size
    for (int i = 0; i <= N; i++) {
        for (int j = 0; j <= N; j++) {
            element_D[i][j] *= 2.0 / dx;
        }
    }
    
    // Assemble global operators
    int global_size = Ncell * N + 1;
    D1.resize(global_size, std::vector<double>(global_size, 0.0));
    D2.resize(global_size, std::vector<double>(global_size, 0.0));
    
    std::cout << "SEM Global Size: " << global_size << std::endl;
    
    // Fill matrices element by element
    int offset = 0;
    for (int cell = 0; cell < Ncell; cell++) {
        for (int i = 0; i <= N; i++) {
            // Skip the first point of each element except for the first element
            if (i == 0 && cell > 0) continue;
            
            int global_i = offset + i;
            if (global_i >= global_size) {
                std::cerr << "Error: global_i (" << global_i << ") exceeds global_size (" << global_size << ")" << std::endl;
                continue;
            }
            
            for (int j = 0; j <= N; j++) {
                // Skip the first point of each element except for the first element
                if (j == 0 && cell > 0) continue;
                
                int global_j = offset + j;
                if (global_j >= global_size) {
                    std::cerr << "Error: global_j (" << global_j << ") exceeds global_size (" << global_size << ")" << std::endl;
                    continue;
                }
                
                D1[global_i][global_j] = element_D[i][j];
            }
        }
        
        // Adjust offset for the next element (skip shared node)
        offset += N;
    }
    
    // Compute D2 = D1 * D1
    for (int i = 0; i < global_size; i++) {
        for (int j = 0; j < global_size; j++) {
            D2[i][j] = 0.0;
            for (int k = 0; k < global_size; k++) {
                D2[i][j] += D1[i][k] * D1[k][j];
            }
        }
    }
    
    std::cout << "SEM Operators Generated: points.size()=" << points.size() 
              << ", D1.size()=" << D1.size() 
              << ", D2.size()=" << D2.size() << std::endl;
}

// Generate Chebyshev points and differentiation matrices
void generate_chebyshev_operators(int N, double L,
                                 std::vector<double>& points,
                                 std::vector<std::vector<double>>& D1,
                                 std::vector<std::vector<double>>& D2) {
    points.resize(N + 1);
    D1.resize(N + 1, std::vector<double>(N + 1, 0.0));
    D2.resize(N + 1, std::vector<double>(N + 1, 0.0));
    
    // Generate Chebyshev points
    for (int i = 0; i <= N; i++) {
        points[i] = -L * std::cos(M_PI * i / N);
    }
    
    // Compute first derivative matrix
    for (int i = 0; i <= N; i++) {
        for (int j = 0; j <= N; j++) {
            if (i == j) {
                if (i == 0) {
                    D1[i][j] = (2.0 * N * N + 1.0) / 6.0;
                } else if (i == N) {
                    D1[i][j] = -(2.0 * N * N + 1.0) / 6.0;
                } else {
                    D1[i][j] = -points[i] / (2.0 * (1.0 - points[i] * points[i]));
                }
            } else {
                double c_i = (i == 0 || i == N) ? 2.0 : 1.0;
                double c_j = (j == 0 || j == N) ? 2.0 : 1.0;
                D1[i][j] = c_i / c_j * std::pow(-1.0, i + j) / (points[i] - points[j]);
            }
        }
    }
    
    // Compute second derivative matrix (D2 = D1 * D1)
    for (int i = 0; i <= N; i++) {
        for (int j = 0; j <= N; j++) {
            D2[i][j] = 0.0;
            for (int k = 0; k <= N; k++) {
                D2[i][j] += D1[i][k] * D1[k][j];
            }
        }
    }
}

void communicate_boundaries(float* data, const LocalGridInfo& info) {
    // Calculate dimensions
    int imt = info.local_imt;
    int jmt = info.local_jmt;
    int kmt = info.local_kmt;
    int padded_imt = imt + 2;
    int padded_jmt = jmt + 2;
    int padded_kmt = kmt + 2;

    // Create MPI requests and status arrays
    MPI_Request requests[6];
    MPI_Status statuses[6];
    int request_count = 0;

    // Send and receive data in each direction
    // North-South
    if (info.north_rank != MPI_PROC_NULL) {
        // Send north boundary
        MPI_Isend(data + padded_imt, padded_imt, MPI_FLOAT, info.north_rank, 0, MPI_COMM_WORLD, &requests[request_count++]);
        // Receive south boundary
        MPI_Irecv(data + (jmt + 1) * padded_imt, padded_imt, MPI_FLOAT, info.south_rank, 0, MPI_COMM_WORLD, &requests[request_count++]);
    }

    // East-West
    if (info.east_rank != MPI_PROC_NULL) {
        // Send east boundary
        for (int k = 0; k < padded_kmt; k++) {
            MPI_Isend(data + k * padded_imt * padded_jmt + padded_jmt - 1, padded_jmt, MPI_FLOAT, info.east_rank, 0, MPI_COMM_WORLD, &requests[request_count++]);
        }
        // Receive west boundary
        for (int k = 0; k < padded_kmt; k++) {
            MPI_Irecv(data + k * padded_imt * padded_jmt, padded_jmt, MPI_FLOAT, info.west_rank, 0, MPI_COMM_WORLD, &requests[request_count++]);
        }
    }

    // Wait for all communications to complete
    if (request_count > 0) {
        MPI_Waitall(request_count, requests, statuses);
    }
}

void exchange_cavity_data(float* data1, float* data2,
                        const LocalGridInfo& info1, const LocalGridInfo& info2,
                        const CavityConnection& connection) {
    // Calculate dimensions for each cavity
    int imt1 = info1.local_imt;
    int jmt1 = info1.local_jmt;
    int kmt1 = info1.local_kmt;
    int imt2 = info2.local_imt;
    int jmt2 = info2.local_jmt;
    int kmt2 = info2.local_kmt;

    // Calculate padded dimensions
    int padded_imt1 = imt1 + 2;
    int padded_jmt1 = jmt1 + 2;
    int padded_kmt1 = kmt1 + 2;
    int padded_imt2 = imt2 + 2;
    int padded_jmt2 = jmt2 + 2;
    int padded_kmt2 = kmt2 + 2;

    // Exchange data based on connection type
    switch (connection.type) {
        case CavityConnection::Type::FACE_X: {
            // Exchange data along X faces
            int face1 = connection.face1_index;
            int face2 = connection.face2_index;
            
            // Copy data from cavity1 to cavity2
            for (int k = 0; k < kmt1; k++) {
                for (int j = 0; j < jmt1; j++) {
                    int idx1 = (k + 1) * padded_imt1 * padded_jmt1 + (j + 1) * padded_imt1 + face1;
                    int idx2 = (k + 1) * padded_imt2 * padded_jmt2 + (j + 1) * padded_imt2 + face2;
                    data2[idx2] = data1[idx1];
                }
            }
            break;
        }
        case CavityConnection::Type::FACE_Y: {
            // Exchange data along Y faces
            int face1 = connection.face1_index;
            int face2 = connection.face2_index;
            
            // Copy data from cavity1 to cavity2
            for (int k = 0; k < kmt1; k++) {
                for (int i = 0; i < imt1; i++) {
                    int idx1 = (k + 1) * padded_imt1 * padded_jmt1 + face1 * padded_imt1 + (i + 1);
                    int idx2 = (k + 1) * padded_imt2 * padded_jmt2 + face2 * padded_imt2 + (i + 1);
                    data2[idx2] = data1[idx1];
                }
            }
            break;
        }
        case CavityConnection::Type::FACE_Z: {
            // Exchange data along Z faces
            int face1 = connection.face1_index;
            int face2 = connection.face2_index;
            
            // Copy data from cavity1 to cavity2
            for (int j = 0; j < jmt1; j++) {
                for (int i = 0; i < imt1; i++) {
                    int idx1 = face1 * padded_imt1 * padded_jmt1 + (j + 1) * padded_imt1 + (i + 1);
                    int idx2 = face2 * padded_imt2 * padded_jmt2 + (j + 1) * padded_imt2 + (i + 1);
                    data2[idx2] = data1[idx1];
                }
            }
            break;
        }
    }
}

void communicate_boundaries_multi_cavity(float* data, const MultiCavityGridParams& grid) {
    // Communicate boundaries for each cavity
    for (int cavity_id = 0; cavity_id < grid.cavities.size(); cavity_id++) {
        const auto& cavity = grid.cavities[cavity_id];
        const auto& local_info = cavity.local_info;
        
        // Calculate offsets for this cavity's data in the global array
        int offset = cavity.data_offset;
        float* cavity_data = data + offset;
        
        // Communicate boundaries for this cavity
        communicate_boundaries(cavity_data, local_info);
    }
    
    // Handle inter-cavity connections
    for (const auto& connection : grid.connections) {
        // Get cavity information
        const auto& cavity1 = grid.cavities[connection.cavity1_id];
        const auto& cavity2 = grid.cavities[connection.cavity2_id];
        
        // Calculate data offsets
        float* data1 = data + cavity1.data_offset;
        float* data2 = data + cavity2.data_offset;
        
        // Exchange data between cavities based on connection type
        exchange_cavity_data(data1, data2, cavity1.local_info, cavity2.local_info, connection);
    }
}

} // namespace TensorMatMultiCavity 